from typing import List

import aiohttp


class RerankerServer:
    """
    https://cloud.baidu.com/doc/qianfan-api/s/2m7u4zt74
    """
    def __init__(self, config: dict = None, apikey: dict = None):
        apikey = apikey if apikey else {}
        self.model: str = ""
        self.base_url: str = "https://qianfan.baidubce.com/v2"
        self.apis_key: str = apikey.get("api_key") or ""
        self.headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.apis_key}"
        }

        if config:
            if config.get("agency_api"):
                self.base_url = config.get("agency_api").rstrip().rstrip("/")
            if config.get("model"):
                self.model = config.get("model")

        if not self.model:
            raise Exception("No rerank model")

    async def send(self, query: str, documents: List[str]):
        async with aiohttp.ClientSession() as session:
            payload = {
                "model": self.model,
                "query": query,
                "documents": documents
            }

            try:
                async with session.post(
                    url=f"{self.base_url}/rerank",
                    headers=self.headers,
                    json=payload
                ) as response:
                    if response.status == 200:
                        return await response.json()
                    else:
                        error_text = await response.text()
                        raise Exception(f"API request failed with status {response.status}: {error_text}")
            except Exception as e:
                raise Exception(f"Error making request: {str(e)}")
